#!/usr/bin/env python

#tree_comparison_api.py

"""
API for running tree comparisons programs

Revision History:

6-14-04: Cathy Lozupone: started writing
7-14-06: Cathy Lozupone: fixed some tests. Made to inherit from object and so
that the EnvsNode can be replaced with another tree object type.
"""

from __future__ import division
from old_cogent.base.tree import PhyloNode
from old_cogent.parse.tree import DndParser
from old_cogent.parse.record import RecordError
from copy import deepcopy
from random import shuffle
from sets import Set
import re
from NexusParser import parse_nexus_tree, parse_PAUP_log, \
     split_tree_info, parse_trans_table, parse_dnd, get_BL_table, parse_taxa, \
     find_fields
from EnvsNode import EnvsNode
from EnvsNode import parse_newick, make_env_dict_text, \
        make_env_dict_text_Weight, get_sat_curve_data_tree, \
        generate_sat_curve_output, get_sat_curve_data_envs
from Numeric import nonzero

class TreeAnalysis(object):
    """session for doing an analysis on a tree
    
Class: TreeAnalysis(interactive=True)

    Usage: session = TreeAnalysis()
        Starts a new session called 'session'.
        Interactive indicates whether the session will be run on the command
            line (True) or from a script (False). By default is True.

    Properties:
        Tree: tree on which all called stats functions are run
        LastSession: saves the current TreeAnalysis object whenever 
            the tree is changed
        DataInfo: maps the name of a node (node.Data) to the node object
        Output: records the results of analyses done during a session as a list
        SessionLog: records the steps performed during a session as a list

    Methods:

    To get detailed info on any method type:
    print TreeAnalysis.method_name.__doc__


    ***Loading a Tree:***
    
    loadTreeFromFile(self, file_path, format, tree_name=None)
            loads a tree from a Nexus or Newick from Arb formatted file
    loadTreeFromString(self, input)
            loads a tree from a string or list of strings


    ***Setting the Branch Lengths:***

    setBlFromFile(self, file_path, format='NexLog')
            sets node BranchLengths from info in a file
    setBlToValue(self, value=1.0)
            sets the BranchLengths property of all nodes to value


    ***Loading Environment Information***

    loadEnvs(self, format, file_path=None):
            sets Envs property of nodes


    ***Modifying Tree Contents (taxa represented by tree)***

    pruneTree(self):
        removes terminal descendants that have no environment assigned
    removeNodeByName(name):
        removes a node from the tree identified by name (Data property)
    removeNodeByLCA(self, first_name, second_name):
        removes a node identified as Last Common Ancestor of 2 nodes
    isolateNodeByName(self, name):
        sets the Tree to a node identified by name
    isolateNodeByLCA(self, first_name, second_name):
        sets the Tree to the Last Common Ancestor of 2 nodes
    restoreLastSession(self):
        sets the current TreeAnalysis session to LastSession


    ***Modify Node Names (Data property of nodes in a tree)***

    nameUnnamedNodes(self):
        assigns an arbitrary name to unnamed nodes
    addEnvsToName(self):
        appends the envs to the node name (Data property of node)


    ***Statistical analysis***

    phylogeneticTestP(self, p_output='Tree', pop_size):
        generates p-values for the Tree using the Phylogenetic (P) test
    uniFracP(self, p_output='Tree', pop_size=1000, Weight=False):
        generates p-values for the Tree using the UniFrac metric
    makeEnvDistanceMatrix(self, Weight=False, pad=None, norm=False):
        generates a Distance Matrix for environments in Tree using UniFrac
    clusterEnvs(self, Weight=False, norm=False):
        uses UPGMA to cluster the environments in a tree based on UniFrac
    UniFracPCA(self, Weight=False, norm=False):
        performs principal coordinates analysis on self.Tree using UniFrac
    PD_rarefaction(self, file_path_base='PD_rarefaction_output', env_list=None
        num_reps = 50, stride=1)
        Creates PD rarefaction curves
    G_rarefaction(self, file_path_base='G_rarefaction_output', env_list=None
        num_reps = 50, stride=1)
        Creates G rarefaction curves
        """

    def __init__(self, interactive=True):
        """set session variables"""
        #Tree is the tree on which all called stats functions are run
        self.Tree = None
        #DataInfo maps the data property to the node
        #it is used for easy lookups and reset whenever data is set
        self.DataInfo = {}

        #LastSession is saved whenever changes are made to Tree
        self.LastSession = None
        
        #determines whether will be run with a script or interactively
        self.interactive = interactive

        #output string for session appended to whenever a result is
        #generated and used to print the results to a file when executing
        #in jobs in batch
        self.Output = []
        self.SessionLog = []

        self.NodeType=EnvsNode

    def loadTreeFromUploadFile(self, lines, format, tree_name=None,
                    assign_branch_len=True):
        """populates Tree with a tree object made from info in a upload 

        Usage: self.loadTreeFromFile(file_path, format, tree_name=None)

        file_path is the path to the tree file

        format is the format of the tree file:
            Nex = Nexus format
            NwA = Newick from Arb
    
        tree_name = the name of the tree to set as the active tree:
            use for files with multiple trees loaded in them

        """
        #clear previous Tree and DataInfo, and set LastSession 
        self._clear_tree()

        #open the file

        if lines:
            #get dnd and trans_table if present
            dnd, trans_table = self._getTreeInfo(lines, format)

            if dnd:
                #set Tree using info in the dendrogram
                self._setTree(dnd, tree_name)


            #translate names with trans_table if it is present
            if trans_table:
                self._translate_node_names(trans_table)

        self._report_tree_info()

        # assign branch length of 1.0 if no branch length included
        if assign_branch_len:
            tree = self.Tree
            bltotal = sum([node.BranchLength for node in tree.traverse() \
                    if node.BranchLength is not None])
            if bltotal == 0:
                for t in tree.traverse():
                    if t.BranchLength is None:
                        t.BranchLength = 1.0
            tree.BranchLength = 0.0
        #SessionLogStr = "Tree Loaded from File: " + file_path + '\n' + \
        #                   str(self.Tree)
        #self.SessionLog.append(SessionLogStr)
        
    def loadEnvsFromUploadFile(self, lines, format='TabDel'):
        """sets Envs property of nodes

        Usage: self.loadEnvsFromUploadFile(lines, format)

        format = format of envs information
            TabDel = tab delimited file. Sequence Name and environments in 
                which that sequence was found separated by tabs.
            TabDelWeight = tab delimited file with Abundance information for                   sequences (e.g. RFLP data). Sequence name, environment name, 
               and number of times the sequence was observed separated by
               tabs. Sequences that were found in more than one environment
               should have a separate line for each environment.

        """
        env_dict = {}
        if self._tree_check():
            if not lines:
                if self.interactive:
                    print "No environment data." 
                else:
                    raise ValueError, "No environment data."
            if not (format == 'TabDel' or format == 'TabDelWeight'):
                if self.interactive:
                    print "Not a valid format code"
                else:
                    raise ValueError, "Not a valid environment format."
            else: 
                if format == 'TabDel':
                    env_dict, all_envs = make_env_dict_text(lines)
                elif format == 'TabDelWeight':
                     env_dict, all_envs = make_env_dict_text_Weight(lines)
            if env_dict:
                self.Tree.assignTerminalEnvs(env_dict, all_envs)
                self._updateDataInfo()
                if self.interactive:
                    print "environment information loaded. All envs =  "
                    print self.Tree.EnvsNames
                    print '\n'
                    a = self._find_unassigned_nodes()

    def loadTreeFromFile(self, file_path, format, tree_name=None):
        """populates Tree with a tree object made from info in a file

        Usage: self.loadTreeFromFile(file_path, format, tree_name=None)

        file_path is the path to the tree file

        format is the format of the tree file:
            Nex = Nexus format
            NwA = Newick from Arb
    
        tree_name = the name of the tree to set as the active tree:
            use for files with multiple trees loaded in them

        """
        #clear previous Tree and DataInfo, and set LastSession 
        self._clear_tree()

        #open the file
        lines = self._open_file(file_path)

        if lines:
            #get dnd and trans_table if present
            dnd, trans_table = self._getTreeInfo(lines, format)

            #translate names with trans_table if it is present
            if dnd:
                #set Tree using info in the dendrogram
                self._setTree(dnd, tree_name)

            if trans_table:
                self._translate_node_names(trans_table)

        self._report_tree_info()
        SessionLogStr = "Tree Loaded from File: " + file_path + '\n' + \
                           str(self.Tree)
        self.SessionLog.append(SessionLogStr)
        
    def _open_file(self, file_path):
        """opens a file"""
        try:
            lines = open(file_path)
            return lines
        except IOError, error_message:
            if self.interactive:
                print error_message
            else:
                raise IOError, error_message

    def _clear_tree(self):
        """clears the current active tree and associated data"""

        self._setLastSession() #user can go back to previous session

        self.Tree = None
        self.DataInfo = {}

    def _setLastSession(self):
        """sets LastSession variable to the current session"""
        
        self.LastSession = deepcopy(self)

    def _getTreeInfo(self, lines, format):
        """returns a dendrogram and trans_table (if present) from a file

        The trans_table maps taxa number to names in Nexus files"""
        dnd = None
        trans_table = None
    
        if format == "NwA":
            try:
                dnd = parse_newick(lines)
            except RecordError:
                error_message = "not a valid Newick file."
                if self.interactive:
                    print error_message
                else:
                    raise RecordError, error_message

        elif format == "Nex":
            try:
                trans_table, dnd = parse_nexus_tree(lines)
            except RecordError:
                error_message = "not a valid Nexus file."
                if self.interactive:
                    print error_message
                else:
                    raise RecordError, error_message

        else:
            if self.interactive:
                print "Not a recognized format code."
            else:
                raise ValueError, "not a recognized format code."

        if dnd == None:
            error_message = "No trees were successfully imported from the file."
            if self.interactive:
                print error_message
            else:
                raise RecordError, error_message

        return dnd, trans_table

    def _setTree(self, dnd, tree_name=None):
        """uses the info in dnd to set Tree

        dnd maps tree name to a list of strings for a Newick dendrogram.
            It can be generated with _getTreeInfo.

        tree_name allows the user to specify which tree to set as active
            when dnd has multiple trees in (e.g. this could happen if
            data is loaded in from a Nexus file with multiple trees in it)
        """
        #if there were no trees give an error message
        if not dnd:
            if self.interactive:
                print "Tree info not available"
            else:
                raise ValueError,  "Tree info not available"
        
        #if there is only one tree set it to the active tree
        elif len(dnd) == 1:
            dnd_key = dnd.keys()[0]
            dnd = dnd[dnd_key]
            self._loadTree(dnd)

        #if there are multiple trees get active tree from tree_name
        else:
            self._multipleTreeHandler(dnd, tree_name)
    
        #if for any reason self.Tree is still none, error
        if not self.Tree:
            if self.interactive:
                print "Tree was not successfully loaded"
            else:
                raise ValueError, "Tree was not successfully loaded"

    def _loadTree(self, dnd):
        """runs DndParser on a dnd, updates DataInfo and checks for errors"""

        try:
            self.Tree = DndParser(dnd, constructor = self.NodeType)

        except RecordError, e:
            if self.interactive:
                print e
            else:
                raise RecordError, e

        if self.Tree:
            #populate DataInfo with info from newly loaded tree
            self._updateDataInfo()
        
    def _updateDataInfo(self):
        """updates dataInfo property: maps node.Data to node object"""

        self.DataInfo = {}
        for node in self.Tree.traverse():
            if node.Data:
                self.DataInfo[node.Data] = node
            else:
                pass
    
    def _multipleTreeHandler(self, dnd, tree_name=None):
        """deals with dnd that have more than one tree

        e.g. generated by loading a Nexus file that had more than one tree
        """
        
        keys = [key for key in dnd.keys()]
        keys = ','.join(keys)
        if tree_name:
            try:
                self._loadTree(dnd)
            except KeyError:
                if self.interactive:
                    print "invalid tree_name: valid values are: " % keys
                else:
                    raise KeyError, \
                        "invalid tree_name: valid values are: %s" % keys
        else:
            if self.interactive:
                print \
                "multiple trees loaded: assign tree with tree_name =: %s" \
                    % keys
            else:
                raise ValueError,\
                "multiple trees loaded: assign tree with tree_name =: %s" \
                % keys

    def loadTreeFromString(self, input):
        """loads a tree from a string or list of strings for a Newick dnd

        Usage: self.loadTreeFromString(input)

        input is a tree that is described in Newick format, input either as
            a string or a list of strings
        """
        #clear previous Tree and DataInfo, and set LastSession
        self._clear_tree()
        if input == str(input):
            dnd = {'tree_1': [input]}
        elif input == list(input):
            dnd = {'tree_1': input}

        self._setTree(dnd)
        self._report_tree_info()

        SessionLogStr = "Tree Loaded from String: \n"  + str(self.Tree)
        self.SessionLog.append(SessionLogStr)

    def setBlToValue(self, value=1.0):
        """sets the BranchLengths property of all nodes to value

        Usage: self.SetBlToValue(value)

        value: value to which branch lengths for all nodes in the tree
            will be set (except for the root for which the branch length
            will be set to None
        """
        if self._tree_check():
            self.Tree.setBranchLength(value)
            self._updateDataInfo()

            SessionLogStr = "All Branch Lengths set to " + str(value)
            self.SessionLog.append(SessionLogStr)

    def setBlFromFile(self, file_path, format='NexLog'):
        """sets node BranchLengths from info in a file

        Usage: self.setBlFromFile(file_path, format='NexLog')

        file_path is the path to the branch length file

        format is the format of the branch length file:
            NexLog = Nexus log file format
        """
        #open the file
        lines = self._open_file(file_path)
        if lines:
            self._set_bl_from_file_handle(lines, format)
            SessionLogStr = "Branch Lengths set from file: " + file_path +\
                           '\n', str(self.Tree)
            self.SessionLog.append(SessionLogStr)

    def _set_bl_from_file_handle(self, lines, format='NexLog'):
        """sets node BranchLengths from info in lines

        lines is an open file object

        format is the format of the branch length file:
            NexLog = Nexus log file format
        """
        if self._tree_check():
            if format == 'NexLog':
                self._set_bl_from_nexus(lines)
            else:
                if self.interactive:
                    print "Not a recognized format code"
                else:
                    raise ValueError, "Not a recognized format code"
                    
            self._report_tree_info()
        
    def _set_bl_from_nexus(self, lines):
        """sets node BranchLengths from info in a Nexus log file"""
        BL_info = parse_PAUP_log(lines)
        self.Tree.assignBranchLength(BL_info)
        self._updateDataInfo()

    def loadEnvs(self, format='TabDel', file_path=None):
        """sets Envs property of nodes

        Usage: self.loadEnvs(format, file_path)

        format = format of envs information
            TabDel = tab delimited file. Sequence Name and environments in 
                which that sequence was found separated by tabs.
            TabDelWeight = tab delimited file with Abundance information for                   sequences (e.g. RFLP data). Sequence name, environment name, 
               and number of times the sequence was observed separated by
               tabs. Sequences that were found in more than one environment
               should have a separate line for each environment.

        file_path = path to file information
            
        """
        env_dict = {}
        if self._tree_check():
            if format == 'TabDel' or format == 'TabDelWeight':
                if not file_path:
                    error_message = "a file_path is needed for this option"
                    if self.interactive:
                        print error_message
                    else:
                        raise ValueError, error_message
                else:
                    lines = self._open_file(file_path)
                    if format == 'TabDel':
                        if lines:
                            env_dict, all_envs = make_env_dict_text(lines)
                    elif format == 'TabDelWeight':
                        if lines:
                            env_dict, all_envs = \
                                    make_env_dict_text_Weight(lines)
            else:
                if self.interactive:
                    print "Not a valid format code"
                else:
                    raise ValueError, "Not a valid format code"
            
            if env_dict:
                self.Tree.assignTerminalEnvs(env_dict, all_envs)
                self._updateDataInfo()
                if self.interactive:
                    print "environment information loaded. All envs =  "
                    print self.Tree.EnvsNames
                    print '\n'
                    a = self._find_unassigned_nodes()

    def _find_unassigned_nodes(self):
        """Finds terminal nodes that have no Envs assigned"""
        unassignedTerminalEnvs = []
        for i in self.DataInfo:
            if not self.DataInfo[i].Children:
                if not self.DataInfo[i].Envs:
                    unassignedTerminalEnvs.append(self.DataInfo[i])
        if unassignedTerminalEnvs:
            if self.interactive:
                print "The following taxa were not assigned to an Environment:"
                for node in unassignedTerminalEnvs:
                    print node.Data, ","
                output_str = "If you are going to evaluate the tree using" +\
                             " significance tests (e.g. self.phylogenetic" +\
                             "TestP or uniFracP), you " +\
                             "should first remove these taxa from the tree" +\
                             " with self.pruneTree"
                print output_str
            return unassignedTerminalEnvs

    def pruneTree(self):
        """Removes terminal descendants that have no environment assigned

        Usage: self.pruneTree()

        Having terminal descendants with no assigned environments
        will make the results inaccurate for the p (parsimony)- test
        but not the d (distance)- test. This test must thus be run
        before doing a p-test
        """
        
        if self._tree_check() and self._env_check():
            self.LastSession = deepcopy(self)
            unassignedTerminalEnvs = self.Tree.returnEmptyTermDecs()
            for node in unassignedTerminalEnvs:
                self.Tree.removeNodeAndOnlyChildParents(node)
            self._updateDataInfo()

            self._report_tree_info()

    def removeNodeByName(self, name):
        """removes a node from the tree that is identified by Name

        Usage: self.removeNodeByName(name)

        name = the Data property of a node
        """
        if self._node_name_check(name):
            self._setLastSession()
            self.Tree.getSubtreeByData(self.DataInfo[name])
            self._updateDataInfo()
            self._report_tree_info()

    def removeNodeByLCA(self, first_name, second_name):
        """removes a node identified as Last Common Ancestor of 2 nodes

        Usage: self.removeNodeByLCA(first_name, second_name)

        name = the Data property of a node
        """
        if self._node_name_check(first_name) and \
                self._node_name_check(second_name):
            self._setLastSession()
            node = self.DataInfo[first_name].\
                   lastCommonAncestor(self.DataInfo[second_name])
            self.Tree.getSubtreeByData(node)
            self._updateDataInfo()
            self._report_tree_info()

    def isolateNodeByName(self, name):
        """sets the Active Tree to a node identified by name

        Usage: self.isolateNodeByName(name)

        name = the Data property of a node
        """
        if self._node_name_check(name):
            self._setLastSession()
            self.Tree = self.Tree.getSubtreeByData(self.DataInfo[name])
            self._updateDataInfo()
            self._report_tree_info()
        
    def isolateNodeByLCA(self, first_name, second_name):
        """sets the Active Tree to the Last Common Ancestor of 2 nodes

        Usage: self.isolateNodeByLCA(first_name, second_name)

        name = the Data property of a node
        """
        if self._node_name_check(first_name) and \
                self._node_name_check(second_name):
            self._setLastSession()
            node = self.DataInfo[first_name].\
                   lastCommonAncestor(self.DataInfo[second_name])
            self.Tree = self.Tree.getSubtreeByData(node)
            self._updateDataInfo()
            self._report_tree_info()

    def restoreLastSession(self):
        """sets the current Session to LastSession (undo)
        
        usage: self = self.restoreLastSession()"""
        self = self.LastSession
        self._report_tree_info()
        return self

    def nameUnnamedNodes(self):
        """assigns an arbitrary name to unnamed nodes"""
        if self._tree_check():
            self.Tree.nameUnnamedNodes()
            self._updateDataInfo()
            self._report_tree_info()
        
    def _translate_node_names(self, trans_table):
        """changes Data property of nodes using a Nexus Translation Table

        The Nexus Translation table is stored in trans_table when the
        tree is loaded from a Nexus file"""

        if self._tree_check():
            self._setLastSession()
            for node in self.Tree.TerminalDescendants:
                    if trans_table.has_key(node.Data):
                        node.Data = trans_table[node.Data]
            self._updateDataInfo()

            self._report_tree_info()

    def addEnvsToName(self):
        """appends the envs to the node name (Data property of node)"""
        if self._tree_check():
            self.Tree.addEnvsToData()
            self._updateDataInfo()

        self._report_tree_info()

    def phylogeneticTestP(self, p_output='Tree', pop_size=1000):
        """generates a p-value for the Tree using Parsimony

        Usage = self.phylogeneticTestP(p_output, pop_size)

        p_output determines the type of p-values that will be output
            Tree = One p-value representing the probability that there are less
                parsimony changes in the tree than expected by chance
            Pairwise = One parsimony based p value for every pair of envs 
            
        pop_size is the number of random trees generated when determining
            the p-value: default = 1000"""

        output_str = self.phylogeneticTestPString(p_output, pop_size)
        if self.interactive:
            print output_str
        self.Output.append(output_str)
        self.SessionLog.append(output_str)

    def phylogeneticTestPString(self, p_output='Tree', pop_size=1000):
        """generates a p-value for the Tree using Parsimony

        Usage = self.phylogeneticTestP(p_output, pop_size)

        p_output determines the type of p-values that will be output
            Tree = One p-value representing the probability that there are less
                parsimony changes in the tree than expected by chance
            Pairwise = One parsimony based p value for every pair of envs 
            
        pop_size is the number of random trees generated when determining
            the p-value: default = 1000
        
        Returns output string. phylogeneticTestP calls this
            """

        if self._tree_check() and self._env_check():
            if self._all_env_check():
                if p_output == 'Tree':
                    p_vals = self.Tree.phylogeneticTestTree(int(pop_size))
                    return "Parsimony P Value(s) - " +  p_vals
                elif p_output == 'Pairwise':
                    result = self.Tree.phylogeneticTestPairwise(pop_size)
                    return 'Pairwise p-values\n' + result.toDelimited()
                else:
                    if self.interactive:
                        print "not a valid analysis type(p_output)"
                    else:
                        raise ValueError, "Not a valid analysis type(p_output)"
            
    def uniFracP(self, p_output='Tree', pop_size=1000, Weight=False):
        """generates p-values for the Active Tree using Distance

        Usage = self.uniFracP(p_output, pop_size, Weight)

        p_output determines the type of p-values that will be output
            Tree = One p-value representing the probability that there are
                more unique branch lengths in the whole tree than expected
                by chance
            Envs = One p-value for each environment in the tree representing
                the probability that each environment has more unique branch
                length than expected by chance. 
            Pairwise = One p-value for each pair of environments in the tree

        pop_size is the number of random trees generated when determining
            the p-value: default = 1000

        Weight: Only relevant is p_output is Pairwise:
            If True, takes abundance information into account,
            if False, does not.
        """
        if self._tree_check() and self._env_check():
            if self._bl_check():
                if p_output == 'Tree':
                    output_str = self._unifrac_tree(int(pop_size))
                elif p_output == 'Envs':
                    output_str = self._unifrac_envs(int(pop_size))
                elif p_output == 'Pairwise':
                    output_str = self._unifrac_pairwise(pop_size, Weight)
                else:
                    if self.interactive:
                        print "Not a valid analysis type (p_output)"
                    else:
                        raise ValueError, "Not a valid analysis type (p_output)"
        if self.interactive:
            print output_str

        self.Output.append(output_str)
        self.SessionLog.append(output_str)

    def uniFracPString(self, p_output='Tree', pop_size=1000, Weight=False):
        """generates p-values for the Active Tree using Distance

        Usage = self.uniFracP(p_output, pop_size)

        p_output determines the type of p-values that will be output
            Tree = One p-value representing the probability that there are
                more unique branch lengths in the whole tree than expected
                by chance
            Envs = One p-value for each environment in the tree representing
                the probability that each environment has more unique branch
                length than expected by chance. 
            Pairwise = One p-value for each pair of environments in the tree

        pop_size is the number of random trees generated when determining
            the p-value: default = 1000

        Weight: Only relevant is p_output is Pairwise:
            If True, takes abundance information into account,
            if False, does not.
        """
        if Weight and p_output != 'Pairwise':
            raise NotImplemented
            
        if self._tree_check() and self._env_check():
            if self._bl_check():
                if p_output == 'Tree':
                    output_str = self._unifrac_tree(int(pop_size))
                elif p_output == 'Envs':
                    output_str = self._unifrac_envs(int(pop_size))
                elif p_output == 'Pairwise':
                    return self.Tree.uniFracPairwise(pop_size, Weight)
                else:
                    if self.interactive:
                        print "Not a valid analysis type (p_output)"
                    else:
                        raise ValueError, "Not a valid analysis type (p_output)"
        return output_str


    def _unifrac_tree(self, pop_size):
        """generates one p-value for the tree using Distance"""
        p_val = self.Tree.uniFracTree(int(pop_size))
        output_str = "Distance P Value(s) - " + str(p_val)
        return output_str
                           
    def _unifrac_envs(self, pop_size):
        """generates one p-value for each env in the tree using Distance
        """
        p_value, raw_p_value = self.Tree.uniFracTreeEnv(int(pop_size))
        output = []
        output.append("Environmental Distance P Value(s)")
        output.append("Environment\tP Value\tRaw P Value")

        for env in p_value:
            output.append("%s\t%s\t%s" % (env, str(p_value[env]),
                                          str(raw_p_value[env])))

        return '\n'.join(output) 

    def _unifrac_pairwise(self, pop_size, Weight):
        """calcs a p-value for each pair of envs in the tree using UniFrac"""
        result = self.Tree.uniFracPairwise(pop_size, Weight)
        output_str = "Pairwise p-values\n" + result.toDelimited()
        return output_str

    def makeEnvDistanceMatrix(self, Weight=False, pad=None, norm=False):
        """returns a matrix with distances between all of the envs in the tree
        
        Weight: False calculates the distnaces with unweighted Unifrac
            True calculates the distances with weighted UniFrac
        
        pad: the value that empty fields (e.g. the values along the diagonal)
            will be padded with

        norm: only relevant for Weighted Unifrac.
                If True, normalizes the Weighted UniFrac distances by the 
                average distance from the root. If False, does not. 
        """
        env_matrix = self.makeEnvDistanceMatrixRaw(Weight=Weight, pad=pad, \
                        norm=norm)
        if env_matrix:
                output_str = "Environment Distance Matrix\n" + \
                        env_matrix.toDelimited()
                if self.interactive:
                        print output_str
                self.Output.append(output_str)
                self.SessionLog.append(output_str)

    def makeEnvDistanceMatrixRaw(self, Weight=False, pad=None, norm=False):
        """returns a matrix with distances between all of the envs in the tree
        
        pad: the value that empty fields (e.g. the values along the diagonal)
            will be padded with
        Weight: If Weight is False, does not take the abundance information 
                into account
        norm: only relevant for Weighted Unifrac.
                If True, normalizes the Weighted UniFrac distances by the 
                average distance from the root. If False, does not. 
        Returns raw env matrix - used w/web shait - other function should call
        """
        env_matrix = None 
        if self._tree_check() and self._env_check():
            z = self._bl_check()
            if z:
                env_matrix = self.Tree.makeEnvMatrix(default=pad, Weight=Weight,\
                        norm=norm)
        return env_matrix

    def clusterEnvs(self, Weight=False, norm=False):
        """returns the Newick string for a tree clustering the environments

        Usage: TreeAnalysis.clusterEnvs(Weight=False)
        
        clusters using UPGMA
        
        If Weight is False, does not take the abundance information into account
        if True, the branch length for a node is weighted based on the relative
            abundances of the descendants of that node among environments
        
        if norm is True, it normalizes the weighted UniFrac values by the
            average root to tip distance. If norm is False it does not.
        """
        if self._tree_check() and self._env_check():
            if self._bl_check():
                env_tree = self.Tree.clusterEnvsUPGMA(Weight, norm=norm)
                output_str = "Environment tree:\n" + env_tree.__str__() + ';'
                if self.interactive:
                    print output_str
                self.Output.append(output_str)
                self.SessionLog.append(output_str)
                return env_tree
        
    def jackClusterOrigEnvs(self, dnd): 
        """
        Cluster environments and assign envs to internal nodes

        Used to compare counts w/jackknifing approach

        E.g. call cluster.assignJackknifeCounts(res)

        where res is the result of ta.Tree.evaluateBalancedUPGMA(...)
        """
        cluster = DndParser(dnd, constructor=self.NodeType)

        cluster.nameUnnamedNodes()
        clusterArray, node_list = cluster.makeTreeArray()
        node_map = dict(zip([x.Data for x in node_list], clusterArray))

        # assign envs to cluster internal nodes
        for t in cluster.traverse():
            if t.Children:
                t.Envs = node_map[t.Data]
        return cluster

    def nodeSigTest(self, algorithm='UniFrac', max_p=0.05, pop_size=1000, \
                    output_path=None):
        """Determines the nodes in the tree that have a significant p-value

        Usage: self.nodeSigTest(algorithm, max_p, pop_size, output_path)

        algorithm is statistical test used to calculate p-value for each node
            UniFrac = unweighted UniFrac p-values calculated
            Phylogenetic = Phylogenetic (P) - test p-values calculated
            UniFracWeighted = Weighted UniFrac p-values are calculated
            ***is UniFrac  by default
        max_p is the maximum p-value for significance, ***default = 0.05
            ***note that this p-value is not Bonferroni corrected - it should
            actually be multiplied by the number of comparisons made
        pop_size is the number of random trees generated to compute p-value
            default = 1000
        output_path: path for file to which output will be saved if running in
            interactive mode. If left as None, output will be printed to
            screen
        """
        sig_nodes = None
        if self._tree_check():
            if algorithm == 'UniFrac':
                if self._env_check() and self._bl_check():
                    sig_nodes, all_nodes = self.Tree._find_sig_nodes_b\
                        (self.NodeType.uniFracTree, float(max_p), int(pop_size))
                    self._report_sig_nodes(sig_nodes, all_nodes, algorithm, \
                            max_p, pop_size, output_path)
            elif algorithm == 'Phylogenetic':
                if self._all_env_check():
                    sig_nodes, all_nodes = self.Tree._find_sig_nodes_b\
                            (self.NodeType.phylogeneticTestTree, float(max_p), \
                            int(pop_size))
                    self._report_sig_nodes(sig_nodes, all_nodes, algorithm, \
                            max_p, pop_size, output_path)
            elif algorithm == 'UniFracWeighted':
                if self._env_check() and self._bl_check():
                    #check that has only 2 envs assigned
                    self.Tree.clearInternalEnvs()
                    self.Tree.assignInternalEnvs()
                    all_envs = [self.Tree.EnvsNames[i] for i in \
                                nonzero(self.Tree.Envs)]
                    if len(all_envs) != 2:
                        raise ValueError, "Tree can only have sequences" +\
                            " assigned to two environments"
                    else:
                        sig_nodes, all_nodes = self.Tree._find_sig_nodes_b\
                            (self.NodeType.uniFracWeightedTree, float(max_p), \
                            int(pop_size))
                        self._report_sig_nodes(sig_nodes, all_nodes, algorithm, \
                                max_p, pop_size, output_path)
            else:
                if self.interactive:
                    print "not a recognized algorithm code"
                else:
                    raise ValueError, "not a recognized algorithm code"                
    def nodeSigTestString(self, algorithm='UniFrac', max_p=0.05, 
                    pop_size=1000):
        """Determines the nodes in the tree that have a significant p-value

        Usage: self.nodeSigTestString(algorithm, max_p, pop_size, output_path)

        algorithm is statistical test used to calculate p-value for each node
            UniFrac = unweighted UniFrac p-values calculated
            Phylogenetic = Phylogenetic (P) - test p-values calculated
            UniFracWeighted = Weighted UniFrac p-values are calculated
            ***is UniFrac  by default
        max_p is the maximum p-value for significance, ***default = 0.05
            ***note that this p-value is not Bonferroni corrected - it should
            actually be multiplied by the number of comparisons made
        pop_size is the number of random trees generated to compute p-value
            default = 1000
       
        returns string - other function should probably call?
        """
        sig_nodes = None
        if self._tree_check():
            if algorithm == 'UniFrac':
                if self._env_check() and self._bl_check():
                    sig_nodes, all_nodes = self.Tree._find_sig_nodes_b\
                        (self.NodeType.uniFracTree, float(max_p), int(pop_size))
                    return self._report_sig_nodes_string(sig_nodes, all_nodes, \
                            algorithm, max_p, pop_size)
            elif algorithm == 'Phylogenetic':
                if self._all_env_check():
                    sig_nodes = self.Tree._find_sig_nodes_b\
                            (self.NodeType.phylogeneticTestTree, float(max_p), \
                            int(pop_size))
                    return self._report_sig_nodes_string(sig_nodes, all_nodes, \
                            algorithm, max_p, pop_size)
            elif algorithm == 'UniFracWeighted':
                if self._env_check() and self._bl_check():
                    #check that has only 2 envs assigned
                    self.Tree.clearInternalEnvs()
                    self.Tree.assignInternalEnvs()
                    all_envs = [self.Tree.EnvsNames[i] for i in \
                                nonzero(self.Tree.Envs)]
                    if len(all_envs) != 2:
                        raise ValueError, "Tree can only have sequences" +\
                            " assigned to two environments"
                    else:
                        sig_nodes, all_nodes = self.Tree._find_sig_nodes_b\
                            (self.NodeType.uniFracWeightedTree, float(max_p), \
                            int(pop_size))
                        return self._report_sig_nodes_string(sig_nodes, \
                                all_nodes, algorithm, max_p,  pop_size)
            else:
                if self.interactive:
                    print "not a recognized algorithm code"
                else:
                    raise ValueError, "not a recognized algorithm code"                
                        
    def _report_sig_nodes(self, sig_nodes, all_nodes, algorithm, max_p, \
            pop_size, output_path):
        """reports results of self.nodeSigTest"""
        #Make output_str
        #Header with information on analysis
        output_str = "Results for nodeSigTest with the " + algorithm +\
                    " algorithm.\nReported significant nodes have a p value "+\
                    "that is <= " + str(max_p) + ".\n" + str(pop_size) +\
                    " random trees were used to calculate the p-value\n\n"
        #Report tree on which the analysis was run
        output_str = output_str + "The analysis was run on the tree:\n\n" +\
                     self.Tree.__str__() + '\n\n'

        #Report p-values for nodes, sig_node groups, and Newick strings
        output_str += self.Tree.generateWithSigOutput(sig_nodes, all_nodes)
                
        if self.interactive:
            if output_path:
                self._print_output_to_file(output_str, output_path)
            else:
                print output_str
                        
        self.Output.append(output_str)
        self.SessionLog.append(output_str)

    def _report_sig_nodes_string(self, sig_nodes, all_nodes, algorithm, max_p, \
            pop_size):
        """reports results of self.nodeSigTestString"""
        #Make output_str
        #Header with information on analysis
        output_str = "Results for nodeSigTest with the " + algorithm +\
                    " algorithm.\nReported significant nodes have a p value "+\
                    "that is <= " + str(max_p) + ".\n" + str(pop_size) +\
                    " random trees were used to calculate the p-value\n\n"
        #Report tree on which the analysis was run
        output_str = output_str + "The analysis was run on the tree:\n\n" +\
                     self.Tree.__str__() + '\n\n'

        #Report p-values for nodes, sig_node groups, and Newick strings
        output_str += self.Tree.generateWithSigOutput(sig_nodes, all_nodes)
                
        return output_str 

    def nodeChiSquare(self, max_p):
        """reports which nodes have a Chi Square value below max_p

        Usage: self.nodeChiSquare(max_p)

        The Chi Square test compares the proportion of taxa belonging to a
        the different environments in a descendants of a node to that of the
        rest of the tree. A significant p value can identify nodes that
        have a concentration of taxa from one environment but not the others.

        max_p: only nodes that have p-values equal or less than max_p will be
        reported
        """
        self.Tree.assignInternalEnvs()
        sig_nodes, all_nodes = self.Tree._find_sig_nodes_a\
                (max_p, self.NodeType.nodeChiSquare)
        output_str = "Results for node Chi Square Test\n\n" + \
        "The following nodes have p_values <= " + str(max_p) + "\n\n" + \
        "Name\tp-value\n"

        for node in sig_nodes:
            output_str = output_str + node[0].Data + '\t' + node[1] + '\n'  
        
        if self.interactive:
            print output_str

        self.Output.append(output_str)

    def UniFracPCA(self, file_path='UniFracPCA_output.txt', Weight=False, \
                    norm=False):
        """Performs Principal Coordinates Analysis based on a UniFrac matrix

        Usage: self.UniFracPCA(file_path, Weight, norm)

        file_path: path to where the output file will be stored. It is
        stored in 'UniFracPCA_output.txt' by default
        
        Weight: if False, branches are not weighted by relative abundance 
        of environments. If True, they are. Weight is set to False by default

        norm: if True, normalizes the Weighted UniFrac values by the average
        distance of the sequences from the root. If False, does not
        normalize.

        Computes a distance matrix for the environments in self.Tree using
        the UniFrac metric and reports the principal coordinates in a 
        tab delimited text file for further analysis in Excel"""
    
        if self._tree_check() and self._env_check() and self._bl_check():
            #generate result
            output_str = self.Tree.clusterEnvsPCA(Weight, norm)
            #output to file
            output_file = open(file_path, 'w')
            output_file.write(output_str)
            output_file.close()
            #write to self.output
            self.Output.append(output_str)
            #interface
            if self.interactive:
                print "UnifracPCA has successfully completed. The results" +\
                        " have been written to file: "
                print file_path

    def PD_rarefaction(self, file_path_base='PD_rarefaction_output_',
                    env_list=None, num_reps=50, stride=1):
        """Creates PD rarefaction curves 

        Usage: self.PD_rarefaction(file_path, env_list, num_reps)

        file_path_base: path to where the output file will be stored.
            The name of the evaluated environment will be appended to 
            this string, followed by .txt
        
        env_list: a list of the environment names to perform rarefaction on;
            If env_list is None, performs rarefaction on all of the 
            environments in the tree.

        num_reps: the number of replicate trials to perform

        stride: The number of sequences removed at a time. For instance
            a stride of 5 for an environment with 15 sequences will return
            the average total branchlength with 15, 10, and 5 sequences in
            the tree.
        
        a tab delimited text file for further analysis in Excel"""
    
        if self._tree_check() and self._env_check() and self._bl_check():
            #generate result
            self.Tree.assignInternalEnvs()
            if not env_list:
                env_list = [self.Tree.EnvsNames[i] for i in \
                        nonzero(self.Tree.Envs)]
            for env in env_list:
                env_tree = deepcopy(self.Tree)
                env_tree.keepEnvTermDecs([env])
                result = get_sat_curve_data_tree(env_tree, num_reps, stride, \
                        Verbose=self.interactive)
                output = generate_sat_curve_output(result)
                output_str = '\n'.join(output)
                #output to file
                outfile_path = file_path_base + env + '.txt'
                output_file = open(outfile_path, 'w')
                output_file.write(output_str)
                output_file.close()
                #write to self.output
                self.Output.append("PD rarefaction results for: " + env)
                self.Output.append(output_str)
                self.Output.append('\n')
                #interface
                if self.interactive:
                    print """PD rarefaction for an environment has completed. 
                    The results have been written to file: """
                    print outfile_path

    def G_rarefaction(self, file_path_base='G_rarefaction_output_',
                    env_list=None, num_reps=50, stride=1):
        """Creates G rarefaction curves 

        Usage: self.G_rarefaction(file_path, env_list, num_reps)

        file_path_base: path to where the output file will be stored.
            The name of the evaluated environment will be appended to 
            this string, followed by .txt
        
        env_list: a list of the environment names to perform rarefaction on;
            If env_list is None, performs rarefaction on all of the 
            environments in the tree.

        num_reps: the number of replicate trials to perform

        stride: The number of sequences removed at a time. For instance
            a stride of 5 for an environment with 15 sequences will return
            the average total branchlength with 15, 10, and 5 sequences in
            the tree.
        
        a tab delimited text file is saved for further analysis in Excel
        
        Warning: This functionality does not work when the same sequence
            if assigned to multiple environments"""
    
        if self._tree_check() and self._env_check() and self._bl_check():
            #generate result
            self.Tree.assignInternalEnvs()
            result, envs = get_sat_curve_data_envs(self.Tree, num_reps, \
                    envs=env_list, stride=stride)
            for index, item in enumerate(result):
                output = generate_sat_curve_output(item)
                output_str = '\n'.join(output)
                #output to file
                outfile_path = file_path_base + envs[index] + '.txt'
                output_file = open(outfile_path, 'w')
                output_file.write(output_str)
                output_file.close()
                #write to self.output
                self.Output.append("G rarefaction results for: " + envs[index])
                self.Output.append(output_str)
                self.Output.append('\n')
                #interface
                if self.interactive:
                    print """G rarefaction for an environment has completed. 
                    The results have been written to file: """
                    print outfile_path

    def _print_output_to_file(self, output_str, file_path = "output.txt"):
        """takes an output file path and string and creates an output file"""
        f = open(file_path, 'w')
        f.write(output_str)
        f.close()

    def _node_name_check(self, name):
        """checks whether a node name is valid"""

        if name not in self.DataInfo:
            if self.interactive:

                    print """***That is not a valid node name***"""
            else:
                    raise KeyError, "Not a valid node name"
        else:
            return True

    def _tree_check(self):
        """Checks if self.Tree has been populated"""
        error_message = "There was a problem loading your tree file. \
                Please check the format and try again."
        if not self.Tree:
            if self.interactive:
                print error_message 
            else:
                raise ValueError, error_message 

        else:
            return True
        
    def _env_check(self):
        """Checks if environment information has been loaded"""
        error_message = \
                "There was a problem loading your environment file. \
                Please check the format and try again."
        Env = False
        for node in self.Tree.traverse():
            if node.Envs:
                Env = True
        if not Env:
            if self.interactive:
                print error_message
            else:
                raise ValueError, error_message
        else:
            return True

    def _all_env_check(self):
        """Checks if all of the taxa have env info assigned

        must be run before executing the parsimony test
        """
        error_message = "Not all taxa have environment info assigned:\n" +\
                    "To remove these taxa from the tree, use self.pruneTree"
        no_env = False
        #sets no_env to True if any terminal 
        for node in self.Tree.TerminalDescendants:
            if node.envEmpty():
                no_env = True
        if no_env:
            if self.interactive:
                print error_message
            else:
                raise ValueError, error_message
        else:
            return True
        
    def _bl_check(self):
        """Checks that BranchLength info is added for the nodes:

        Checks that term dec of tree have BranchLength property"""
        BL = True
        for node in self.Tree.traverse():
            if node.Parent:
                if node.BranchLength == None:
                    BL = False

        error_message = \
                "branch length information is needed for required action"
        if not BL:
            if self.interactive:
                print error_message
            else:
                raise ValueError, error_message
        else:
                return True

    def _report_tree_info(self):
        """prints info on active tree when in interactive mode"""

        if self.interactive:
            print "Tree = ", self.Tree, ';'

            if self.Tree:
                print "\nThere are ", len(self.Tree.TerminalDescendants),\
                      " taxa in the tree"

#=======================================================
if __name__ == '__main__':
    from tree_comparison_api import TreeAnalysis, help

    a = TreeAnalysis()

    print """Welcome to the tree comparison api:

    For a description of the api and available functions type help()

    A tree analysis session 'a' has been started.

    To make another session type:
    session_name = TreeAnalysis()

    To begin, either

    1) load a tree from a Nexus formatted tree file using:
    a.loadTreeFromFile('file_path', 'Nex')

    2) load a tree from a Newick formatted tree file exported from Arb using:
    a.loadTreeFromFile('file_path', 'NwA')

    or 3) load a tree by pasting in a Newick string using:
    a.loadTreeFromString('Newick_string')
    """

            
        



